/**
 * encrypt_plugin.c - 文件加密/解密插件
 * 
 * 此插件演示如何通过钩子框架实现文件加密/解密功能
 */

#define _GNU_SOURCE
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <dlfcn.h>
#include <errno.h>
#include <fcntl.h>
#include <unistd.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <stdarg.h>
#include "../../include/hook_framework.h"

#define ENCRYPT_EXT ".enc"
#define DECRYPT_EXT ".dec"

/* 是否启用加密功能 */
static int g_encryption_enabled = 0;

/* 加密密钥 */
static unsigned char g_encryption_key[16] = {0};

/* 调试模式 */
static int g_debug_mode = 1;

/**
 * 简单XOR加密/解密
 */
static void xor_transform(unsigned char* data, size_t len) {
    for (size_t i = 0; i < len; i++) {
        data[i] ^= g_encryption_key[i % sizeof(g_encryption_key)];
    }
}

/**
 * 检查文件是否需要加密/解密
 * 返回：1-需要加密，-1-需要解密，0-不需要处理
 */
static int check_file_processing(const char* path) {
    if (!path || !g_encryption_enabled) {
        return 0;
    }

    size_t len = strlen(path);
    
    /* 检查是否是.enc文件(需要解密) */
    if (len > 4 && strcmp(path + len - 4, ENCRYPT_EXT) == 0) {
        if (g_debug_mode) fprintf(stderr, "[加密插件] 文件 %s 需要解密\n", path);
        return -1; /* 需要解密 */
    }
    
    /* 检查是否是需要加密的文件类型 */
    const char* extensions[] = {".txt", ".md", ".csv", ".json", ".xml", NULL};
    for (int i = 0; extensions[i] != NULL; i++) {
        size_t ext_len = strlen(extensions[i]);
        if (len > ext_len && strcmp(path + len - ext_len, extensions[i]) == 0) {
            if (g_debug_mode) fprintf(stderr, "[加密插件] 文件 %s 需要加密\n", path);
            return 1; /* 需要加密 */
        }
    }
    
    if (g_debug_mode) fprintf(stderr, "[加密插件] 文件 %s 不需要处理\n", path);
    return 0; /* 不需要处理 */
}

/**
 * 处理文件路径（添加或删除扩展名）
 */
static char* process_path(const char* original_path, int processing_type) {
    if (!original_path) return NULL;
    
    size_t len = strlen(original_path);
    char* new_path = NULL;
    
    if (processing_type == 1) { /* 加密：添加.enc */
        new_path = malloc(len + strlen(ENCRYPT_EXT) + 1);
        if (!new_path) return NULL;
        
        strcpy(new_path, original_path);
        strcat(new_path, ENCRYPT_EXT);
        if (g_debug_mode) fprintf(stderr, "[加密插件] 处理路径: %s -> %s\n", original_path, new_path);
    }
    else if (processing_type == -1) { /* 解密：删除.enc */
        new_path = strdup(original_path);
        if (!new_path) return NULL;
        
        new_path[len - strlen(ENCRYPT_EXT)] = '\0';
        if (g_debug_mode) fprintf(stderr, "[加密插件] 处理路径: %s -> %s\n", original_path, new_path);
    }
    
    return new_path;
}

/* 获取原始函数指针 */
typedef int (*open_func_t)(const char *pathname, int flags, ...);
typedef ssize_t (*read_func_t)(int fd, void *buf, size_t count);
typedef ssize_t (*write_func_t)(int fd, const void *buf, size_t count);

/**
 * open函数的拦截
 */
int open(const char *pathname, int flags, ...) {
    static open_func_t original_open = NULL;
    
    /* 初始化原始函数指针 */
    if (!original_open) {
        original_open = (open_func_t)dlsym(RTLD_NEXT, "open");
        if (!original_open) {
            fprintf(stderr, "无法获取原始open函数: %s\n", dlerror());
            exit(1);
        }
    }
    
    /* 获取mode参数（如果有） */
    mode_t mode = 0;
    if (flags & O_CREAT) {
        va_list args;
        va_start(args, flags);
        mode = va_arg(args, mode_t);
        va_end(args);
    }
    
    if (g_debug_mode) fprintf(stderr, "[加密插件] open调用: %s, flags=%d\n", pathname, flags);
    
    /* 检查文件是否需要加密/解密 */
    int processing_type = check_file_processing(pathname);
    if (processing_type == 0) {
        /* 不需要处理，直接调用原始函数 */
        if (flags & O_CREAT) {
            return original_open(pathname, flags, mode);
        } else {
            return original_open(pathname, flags);
        }
    }

    /* 需要加密或解密 */
    char* processed_path = process_path(pathname, processing_type);
    if (!processed_path) {
        errno = ENOMEM;
        return -1;
    }

    /* 对于加密操作，我们将写入到.enc文件 */
    if (processing_type == 1 && (flags & O_WRONLY || flags & O_RDWR)) {
        if (g_debug_mode) fprintf(stderr, "[加密插件] 加密写入文件: %s -> %s\n", pathname, processed_path);
        int fd;
        if (flags & O_CREAT) {
            fd = original_open(processed_path, flags, mode);
        } else {
            fd = original_open(processed_path, flags);
        }
        free(processed_path);
        return fd;
    }
    /* 对于解密操作，我们将读取.enc文件 */
    else if (processing_type == -1 && (flags & O_RDONLY || flags & O_RDWR)) {
        if (g_debug_mode) fprintf(stderr, "[加密插件] 解密读取文件: %s\n", pathname);
        int fd = original_open(pathname, flags);
        free(processed_path);
        return fd;
    }
    
    /* 其他情况，使用原始路径 */
    free(processed_path);
    if (flags & O_CREAT) {
        return original_open(pathname, flags, mode);
    } else {
        return original_open(pathname, flags);
    }
}

/**
 * read函数的拦截
 */
ssize_t read(int fd, void *buf, size_t count) {
    static read_func_t original_read = NULL;
    
    /* 初始化原始函数指针 */
    if (!original_read) {
        original_read = (read_func_t)dlsym(RTLD_NEXT, "read");
        if (!original_read) {
            fprintf(stderr, "无法获取原始read函数: %s\n", dlerror());
            exit(1);
        }
    }
    
    /* 先调用原始函数 */
    ssize_t bytes_read = original_read(fd, buf, count);
    if (bytes_read <= 0 || !g_encryption_enabled) {
        return bytes_read;
    }

    /* 检查是否需要解密 */
    char fd_path[256] = {0};
    snprintf(fd_path, sizeof(fd_path), "/proc/self/fd/%d", fd);
    
    char file_path[256] = {0};
    ssize_t len = readlink(fd_path, file_path, sizeof(file_path) - 1);
    if (len <= 0) {
        return bytes_read;
    }
    file_path[len] = '\0';
    
    /* 只解密.enc文件 */
    size_t path_len = strlen(file_path);
    if (path_len > 4 && strcmp(file_path + path_len - 4, ENCRYPT_EXT) == 0) {
        /* 进行解密 */
        xor_transform((unsigned char*)buf, bytes_read);
        if (g_debug_mode) fprintf(stderr, "[加密插件] 已解密从%s读取的%zd字节\n", file_path, bytes_read);
    }
    
    return bytes_read;
}

/**
 * write函数的拦截
 */
ssize_t write(int fd, const void *buf, size_t count) {
    static write_func_t original_write = NULL;
    
    /* 初始化原始函数指针 */
    if (!original_write) {
        original_write = (write_func_t)dlsym(RTLD_NEXT, "write");
        if (!original_write) {
            fprintf(stderr, "无法获取原始write函数: %s\n", dlerror());
            exit(1);
        }
    }
    
    if (!g_encryption_enabled) {
        return original_write(fd, buf, count);
    }

    /* 检查是否需要加密 */
    char fd_path[256] = {0};
    snprintf(fd_path, sizeof(fd_path), "/proc/self/fd/%d", fd);
    
    char file_path[256] = {0};
    ssize_t len = readlink(fd_path, file_path, sizeof(file_path) - 1);
    if (len <= 0) {
        return original_write(fd, buf, count);
    }
    file_path[len] = '\0';
    
    if (g_debug_mode) fprintf(stderr, "[加密插件] write调用: fd=%d, 文件=%s, 大小=%zu\n", fd, file_path, count);
    
    /* 只加密写入到.enc文件的数据 */
    size_t path_len = strlen(file_path);
    if (path_len > 4 && strcmp(file_path + path_len - 4, ENCRYPT_EXT) == 0) {
        /* 创建临时缓冲区进行加密 */
        unsigned char* encrypted_buf = malloc(count);
        if (!encrypted_buf) {
            return original_write(fd, buf, count);
        }
        
        /* 复制并加密数据 */
        memcpy(encrypted_buf, buf, count);
        xor_transform(encrypted_buf, count);
        
        /* 写入加密数据 */
        ssize_t result = original_write(fd, encrypted_buf, count);
        free(encrypted_buf);
        
        if (g_debug_mode) fprintf(stderr, "[加密插件] 已加密写入到%s的%zd字节\n", file_path, count);
        return result;
    }
    
    /* 对其他文件使用原始write */
    return original_write(fd, buf, count);
}

/**
 * 加载插件
 */
int plugin_init(void) {
    fprintf(stderr, "[加密插件] 文件加密插件正在初始化\n");
    
    /* 设置加密密钥 */
    const char* key_str = getenv("HOOK_ENCRYPT_KEY");
    if (key_str) {
        size_t key_len = strlen(key_str);
        for (size_t i = 0; i < sizeof(g_encryption_key); i++) {
            g_encryption_key[i] = key_str[i % key_len];
        }
        g_encryption_enabled = 1;
        fprintf(stderr, "[加密插件] 加密功能已启用，使用自定义密钥\n");
    } else {
        /* 使用默认密钥 */
        for (int i = 0; i < sizeof(g_encryption_key); i++) {
            g_encryption_key[i] = (unsigned char)(i + 1);
        }
        g_encryption_enabled = 1;
        fprintf(stderr, "[加密插件] 加密功能已启用(使用默认密钥)\n");
    }
    
    /* 检查是否禁用加密 */
    const char* disable_str = getenv("HOOK_ENCRYPT_DISABLE");
    if (disable_str && (strcmp(disable_str, "1") == 0 || 
                        strcmp(disable_str, "yes") == 0 || 
                        strcmp(disable_str, "true") == 0)) {
        g_encryption_enabled = 0;
        fprintf(stderr, "[加密插件] 加密功能已禁用\n");
    }
    
    /* 检查是否启用调试 */
    const char* debug_str = getenv("HOOK_ENCRYPT_DEBUG");
    if (debug_str && (strcmp(debug_str, "1") == 0 || 
                      strcmp(debug_str, "yes") == 0 || 
                      strcmp(debug_str, "true") == 0)) {
        g_debug_mode = 1;
        fprintf(stderr, "[加密插件] 调试模式已启用\n");
    } else {
        g_debug_mode = 0;
    }
    
    return 0;
}

/**
 * 卸载插件
 */
void plugin_cleanup(void) {
    fprintf(stderr, "[加密插件] 文件加密插件正在关闭\n");
} 